import torch
from torch import nn
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv
from torch_geometric.nn import GATConv, Linear, to_hetero


class Classifier(nn.Module):
    def forward(self, x_A, x_B, index_edge_label):
        edge_feat_A = x_A[index_edge_label[0]]
        edge_feat_B = x_B[index_edge_label[1]]
        return (edge_feat_A * edge_feat_B).sum(dim=-1)


class GNN(nn.Module):
    def __init__(self, hidden_channels):
        super(GNN, self).__init__()
        self.conv1 = SAGEConv(hidden_channels, hidden_channels)
        self.conv2 = SAGEConv(hidden_channels, hidden_channels)

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)
        return x


class Model(nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.gnn = GNN(hidden_channels)
        self.classifier = Classifier()

    def change_model(self, data):
        self.gnn = to_hetero(self.gnn, metadata=data.metadata())
        # pass

    def forward(self, batch_data):
        """
        """
        batch_data0 = batch_data[0].to("cuda:0")
        batch_data1 = batch_data[1].to("cuda:0")
        batch_data2 = batch_data[2].to("cuda:0")

        aoh_label_index = batch_data0[("alarm", "on", "host")].edge_label_index
        # aoh_edge_label = batch_data0[("alarm", "on", "host")].edge_label
        atb_label_index = batch_data1[("alarm", "to", "bussiness_tree")].edge_label_index
        # atb_edge_label = batch_data1[("alarm", "to", "bussiness_tree")].edge_label
        hbb_label_index = batch_data2[("host", "belongsto", "bussiness_tree")].edge_label_index
        # hbb_edge_label = batch_data2[("host", "belongsto", "bussiness_tree")].edge_label

        x_dict0 ={
            "alarm": batch_data0["alarm"].x,
            "host": batch_data0["host"].x,
            "bussiness_tree": batch_data0["bussiness_tree"].x
        }

        out_label_index_dict0 = self.gnn(x_dict0, batch_data0.edge_index_dict)
        out0 = self.classifier(out_label_index_dict0["alarm"], out_label_index_dict0["host"], aoh_label_index)

        x_dict1 = {
            "alarm": batch_data1["alarm"].x,
            "host": batch_data1["host"].x,
            "bussiness_tree": batch_data1["bussiness_tree"].x
        }

        out_label_index_dict1 = self.gnn(x_dict1, batch_data1.edge_index_dict)
        out1 = self.classifier(out_label_index_dict1["alarm"], out_label_index_dict1["host"], atb_label_index)


        x_dict2 = {
            "alarm": batch_data2["alarm"].x,
            "host": batch_data2["host"].x,
            "bussiness_tree": batch_data2["bussiness_tree"].x
        }

        out_label_index_dict2 = self.gnn(x_dict2, batch_data2.edge_index_dict)
        out2 = self.classifier(out_label_index_dict2["alarm"], out_label_index_dict2["host"], hbb_label_index)



        # atb_out = self.classifier(out_label_index_dict["alarm"], out_label_index_dict["bussiness_tree"],
        #                           atb_label_index)
        # # hbb_out = self.classifier(out_label_index_dict["host"], out_label_index_dict["bussiness_tree"], hbb_label_index)
        # # return aoh_out, atb_out, hbb_out

        return out0,out1,out2